{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/RouterQueryEngine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Router Query Engine\n",
    "In this tutorial, we define a custom router query engine that selects one out of several candidate query engines to execute a query."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: This is ONLY necessary in jupyter notebook.\n",
    "# Details: Jupyter runs an event-loop behind the scenes.\n",
    "#          This results in nested event-loops when we start an event-loop to make async queries.\n",
    "#          This is normally not allowed, we use nest_asyncio to allow it for convenience.\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().handlers = []\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SummaryIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    ServiceContext,\n",
    "    StorageContext,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../data/paul_graham\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize service context (set chunk size)\n",
    "service_context = ServiceContext.from_defaults(chunk_size=1024)\n",
    "nodes = service_context.node_parser.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize storage context (by default it's in-memory)\n",
    "storage_context = StorageContext.from_defaults()\n",
    "storage_context.docstore.add_documents(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Summary Index and Vector Index over Same Data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_index = SummaryIndex(nodes, storage_context=storage_context)\n",
    "vector_index = VectorStoreIndex(nodes, storage_context=storage_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Query Engines and Set Metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_query_engine = summary_index.as_query_engine(\n",
    "    response_mode=\"tree_summarize\",\n",
    "    use_async=True,\n",
    ")\n",
    "vector_query_engine = vector_index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.query_engine import QueryEngineTool\n",
    "\n",
    "\n",
    "list_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=list_query_engine,\n",
    "    description=(\n",
    "        \"Useful for summarization questions related to Paul Graham eassy on\"\n",
    "        \" What I Worked On.\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "vector_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=vector_query_engine,\n",
    "    description=(\n",
    "        \"Useful for retrieving specific context from Paul Graham essay on What\"\n",
    "        \" I Worked On.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Router Query Engine\n",
    "\n",
    "There are several selectors available, each with some distinct attributes.\n",
    "\n",
    "The LLM selectors use the LLM to output a JSON that is parsed, and the corresponding indexes are queried.\n",
    "\n",
    "The Pydantic selectors (currently only supported by `gpt-4-0613` and `gpt-3.5-turbo-0613` (the default)) use the OpenAI Function Call API to produce pydantic selection objects, rather than parsing raw JSON.\n",
    "\n",
    "For each type of selector, there is also the option to select 1 index to route to, or multiple.\n",
    "\n",
    "#### PydanticSingleSelector\n",
    "\n",
    "Use the OpenAI Function API to generate/parse pydantic objects under the hood for the router selector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_engine.router_query_engine import RouterQueryEngine\n",
    "from llama_index.selectors.llm_selectors import (\n",
    "    LLMSingleSelector,\n",
    "    LLMMultiSelector,\n",
    ")\n",
    "from llama_index.selectors.pydantic_selectors import (\n",
    "    PydanticMultiSelector,\n",
    "    PydanticSingleSelector,\n",
    ")\n",
    "\n",
    "\n",
    "query_engine = RouterQueryEngine(\n",
    "    selector=PydanticSingleSelector.from_defaults(),\n",
    "    query_engine_tools=[\n",
    "        list_tool,\n",
    "        vector_tool,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the summary of the document?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What did Paul Graham do after RICS?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### LLMSingleSelector\n",
    "\n",
    "Use OpenAI (or any other LLM) to parse generated JSON under the hood to select a sub-index for routing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = RouterQueryEngine(\n",
    "    selector=LLMSingleSelector.from_defaults(),\n",
    "    query_engine_tools=[\n",
    "        list_tool,\n",
    "        vector_tool,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the summary of the document?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What did Paul Graham do after RICS?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] look at selected results\n",
    "print(str(response.metadata[\"selector_result\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PydanticMultiSelector\n",
    "\n",
    "In case you are expecting queries to be routed to multiple indexes, you should use a multi selector. The multi selector sends to query to multiple sub-indexes, and then aggregates all responses using a summary index to form a complete answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SimpleKeywordTableIndex\n",
    "\n",
    "keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)\n",
    "\n",
    "keyword_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=vector_query_engine,\n",
    "    description=(\n",
    "        \"Useful for retrieving specific context using keywords from Paul\"\n",
    "        \" Graham essay on What I Worked On.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = RouterQueryEngine(\n",
    "    selector=PydanticMultiSelector.from_defaults(),\n",
    "    query_engine_tools=[\n",
    "        list_tool,\n",
    "        vector_tool,\n",
    "        keyword_tool,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This query could use either a keyword or vector query engine, so it will combine responses from both\n",
    "response = query_engine.query(\n",
    "    \"What were noteable events and people from the authors time at Interleaf\"\n",
    "    \" and YC?\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] look at selected results\n",
    "print(str(response.metadata[\"selector_result\"]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
